#!/usr/bin/python

from __future__ import division
import arithmetic_parser
from test.test_iterlen import len

def create_function(function_dict):
  """Creates a single function from arithmetic-variety code (for single function)

  Args:
    function_dict: Dictionary holding function information. The keys of the dict
        are 'fname' for function name (should be discarded), 'params' for a list
        of parameters for this function and 'equations' for a list of equations
        this function should solve.
        Each equation is a tupple of the form (dict of parameters with
        coefficients for values, equations solution)
  Returns:
    A python function which accepts the parameters described in the input, solves
        The equations, and returns a dictionary with keys for variables_equations of the
        input, and values for solutions for this keys.
  Raises:
    InvalidInputError: If the equations wouldn't be solvable at runtime.
  """  
  row_canonical_matrix = []
  results_vec = [] 
  if function_dict['equations']:      
      if function_dict['params']:
          variables_equations = [({param: 1}, param) for param in function_dict['params']]
          function_dict['equations'] += variables_equations           
      # create a matrix with parameters' coefficients from equations,
      # a vector from the equations result part and a list of matrix variables
      matrix, results_vec, matrix_vars = create_matrix(function_dict['equations'])     
      num_rows = len(matrix)
      num_cols = len(matrix[0])
      if num_cols > num_rows:
          raise InvalidInputError("More variables than equations, no single solution exists")  
         
      # eliminate matrix rows using Gaussian elimination algorithm
      row_canonical_matrix, results_vec = gauss_elimination(matrix, results_vec)
      # check if matrix is square, if not, check if it contains zero rows
      # that can be ignored
      zero_rows = getZeroRows(row_canonical_matrix)
      zero_cols = getZeroCols(row_canonical_matrix)
      if zero_cols:
          if len(zero_cols) == 1:
              msg = "Equations have a free variable, no single solution exists"
          else:
              msg = "Equations have %s free variables, no single solution exists" % len(zero_cols)
          raise InvalidInputError(msg)
      if len(zero_rows) == num_rows:   #zero matrix
          raise InvalidInputError("All coefficients are zero, no single solution or no solution exists")
      if num_rows - len(zero_rows) != num_cols:
          raise InvalidInputError("Number of equations is different from number of variables, no single solution exists")
      
      if not function_dict['params']:
          for i in xrange(len(results_vec)):
              if type(results_vec[i]) == str:
                  results_vec[i] = eval(results_vec[i])  
          if len(zero_rows) > 0:
              for i in zero_rows:
                  if results_vec[i] != 0:
                      raise InvalidInputError("Linear equations contain equation 0=x, no solution exists") 
                  
                    
  def function_to_return(*args, **kwargs):   
      if args and kwargs:
          raise InvalidInputError("Expect only args or kwargs, but not both")
      elif kwargs:
          args = [kwargs[param] for param in function_dict['params']]
    
      if len(args) != len(function_dict['params']):
          raise TypeError("The number of parameters given doesn't match the required number")
      
      calculated_function_dict = {}
      # create a dictionary with parameters from function_dict
      # and values from input to pass to eval 
      params_dict = {} 
      vector = results_vec[:]         
      if function_dict['equations']:
          if function_dict['params']:
              params = function_dict['params']                
              for param in params:
                  params_dict[param] = args[params.index(param)]                      
              vector = [eval(str(row), params_dict) for row in vector]             
          if num_rows == num_cols:
              calculated_function_dict = getSolution(row_canonical_matrix, vector, matrix_vars)
          else:       #number of zero rows + number of columns = number of rows
             # remove zero rows 
              removeZeroRows(row_canonical_matrix, vector, zero_rows)                  
              calculated_function_dict = getSolution(row_canonical_matrix, vector, matrix_vars)
      else: 
          if function_dict['params']:
              params = function_dict['params']       
              for i in xrange(len(params)):
                  calculated_function_dict[params[i]] = args[i]
      
      return calculated_function_dict
  
  return function_to_return


def create_matrix(equations_list):
    """ Creates a matrix and vector representation from the set of equations
    
    Args:
        equations_list - a list of equations. Each equation is a tuple of a dictionary
                         of variables and their coefficients, and the result of the
                         equation
    Returns:
        A matrix representation of the given equations.
        A result vector containing the equations results
        A sorted list of the matrix's variables
    """
    variables = []
    equations_vars = []    
    # gets all the equations variables
    for equation in equations_list:
        equations_vars += equation[0].keys()       
    # add the equations variables into a list if not already there     
    for var in equations_vars:
        if var not in variables:
            variables.append(var)         
    variables.sort()
    num_cols = len(variables)
    num_rows = len(equations_list)  
    # create a [num_rows]x[num_cols] zero matrix 
    matrix = create_zero_matrix(num_rows, num_cols) 
    # fill in the vector according to the values in equations_list
    results_vec = []
    for row in xrange(num_rows):
        equation = equations_list[row]
        coeff_cols = equation[0]
        result_elem = equation[1]
        for col in xrange(num_cols):
            var = variables[col]
            coeff = coeff_cols.get(var) 
            if coeff is not None:
                matrix[row][col] = coeff
        results_vec.append(result_elem)
        
    return matrix, results_vec, variables

def create_zero_matrix(rows, cols):
    matrix = []   
    for row in xrange(rows):
        matrix.append([0] * cols)
        
    return matrix


def gauss_elimination(matrix, results_vec):
    """ Uses the Gauss Elimination algorithm to solve a set of linear equations
    
    Args:
        A matrix representing the linear equations to solve
        A vector containing the results of the equations
    Returns:
        A matrix in its row canonical form
        A result vector
    """
    num_rows, matrix_num_cols = len(matrix), len(matrix[0])
    i = 0
    k = 0
    while i < num_rows and k < matrix_num_cols:
        # Find pivot in column j, starting in row i
        pivot_index = i
        for j in xrange(i+1, num_rows):
            if abs(matrix[j][k]) > abs(matrix[pivot_index][k]):
                pivot_index = j
        if matrix[pivot_index][k] != 0:
            matrix[i], matrix[pivot_index] = matrix[pivot_index], matrix[i]
            results_vec[i], results_vec[pivot_index] = results_vec[pivot_index], results_vec[i]
            for j in xrange(i+1, num_rows):
                if matrix[j][k] != 0:
                    coeff = matrix[j][k] / matrix[i][k]
                    for x in xrange(k, matrix_num_cols):
                        matrix[j][x] -= matrix[i][x] * coeff
                    results_vec[j] = "(%s - (%s) * %s)" % (results_vec[j], results_vec[i], coeff)
            i += 1
        k += 1
        
    # by now, we are supposed to have an upper triangular matrix,
    # lets turn it into a diagonal matrix
    i = num_rows - 1
    k = matrix_num_cols - 1
    while i >= 0 and k >= 0:
        coeff = matrix[i][k]
        if coeff != 0:
            for j in xrange(0,i):
                results_vec[j] = "(%s - (%s * %s / %s))" % (results_vec[j], results_vec[i], matrix[j][i], coeff)
                for x in xrange(matrix_num_cols-1, i-1, -1):
                    matrix[j][x] -= matrix[i][x] * matrix[j][i] / coeff
            matrix[i][k] /= coeff
            results_vec[i] = "%s / %s" % (results_vec[i], coeff)
            k -= 1
        i -= 1
    
    return matrix, results_vec 
 

def getZeroRows(matrix):
    """ Returns a list with numbers of zero rows
    
    Args:
        matrix - the matrix to check if contains zero rows
    Return:
        A list with the numbers of zero rows
    """
    # search for zero rows
    indices = []
    zero_rows_num = []        
    for i in xrange(len(matrix)):
       if matrix[i][0] == 0:
            indices.append(0)
       else:
            indices.append(1)                
    for i in xrange(len(indices)):
        if indices[i] == 0:
            for j in xrange(len(matrix[0])):
                if matrix[i][j] != 0:
                    break
            else:
                zero_rows_num.append(i)
               
    return zero_rows_num


def getZeroCols(matrix):
    """ Returns a list with numbers of zero columns
    
    Args:
        matrix - the matrix to check if contains zero columns
    Return:
        A list with the numbers of zero columns
    """
    # search for zero columns
    indices = []
    zero_cols_num = []
    for i in xrange(len(matrix[0])):
        if matrix[0][i] == 0:
            indices.append(0)
        else:
            indices.append(1)
    for i in xrange(len(indices)):
        if indices[i] == 0:
            for j in xrange(len(matrix)):
                if matrix[j][i] != 0:
                    break
            else:
                zero_cols_num.append(i)
                
    return zero_cols_num
                
       
def removeZeroRows(matrix, results_vec, indices): 
    """ Removes zero rows from the given matrix
    
    if possible, otherwise throws exception
    Args:
        A matrix in its row canonical form
        A results vector
        indices vector with the number of zero rows
    Raises:
       InvalidInputError - in case there is a zero line in the matrix and 
                           not in the vector in the appropriate place 
    """           
    # "throw" zero rows if results_vec in the same row is 0
    for i in xrange(len(indices)-1, -1, -1):
        if results_vec[indices[i]] == 0:
            matrix.pop(indices[i])
            results_vec.pop(indices[i])
        else:                                
            raise InvalidInputError("Linear equations contain equation 0=x, no solution exists")
       
def getSolution(matrix, vector, variables): 
    """ Returns the solution of the linear equations
    
    Args:
        matrix - the linear equations matrix representation in its row canonical form 
        vector - a vector holding the matrix results
        variables - the linear equations variables
    Return:
        a dictionary with the linear equations solution
    Raises:
        InvalidInputError - in case there is a free variable (zero on the matrix diagonal)
    """    
    solutions_dict = {}
    for i in xrange(len(matrix)):
        if matrix[i][i] != 0:
           solutions_dict[(variables[i])] = vector[i]
        else:
           raise InvalidInputError("Linear equations have a free variable, no single solution exists")
                   
    return solutions_dict

class InvalidInputError(Exception):
    pass            

def compile(code):
  """Compiles given arithmetic variety code.

  Args:
    code: Code as described in the documentation for arithmetic variety
  Return:
    Dictionary mapping name of function in the code to actual python function
    that calculates the result. Result is returned as a dictionary with all
    variables of the code-function (both input and local) as keys.
    None on failure.
  """
  # Dictionary mapping function name to actual function.
  parsed_code = arithmetic_parser.parse(code)
  if parsed_code is None:
    return None
  ret_dict = {}
  for function_dict in parsed_code:
    ret_dict[function_dict['fname']] = create_function(function_dict)
  return ret_dict
  